"""
Whisper实现类
基于OpenAI Whisper的STT实现
"""

import os
from typing import Tuple, Dict, Any

import whisper

# 导入STT实现基类
from .base import STTImplementation
# 导入返回码枚举
from ...constants.return_codes import BaseCodeEnum, STTErrorCode, ReturnCode
# 导入日志工具
from ...utils.logger import get_logger


class WhisperImplementation(STTImplementation):
    """
    基于OpenAI Whisper的STT实现
    """

    # 支持的语言映射字典
    SUPPORTED_LANGUAGES = {
        'zh': '中文',
        'en': '英文',
        'ja': '日文',
        'ko': '韩文',
        'fr': '法文',
        'de': '德文',
        'es': '西班牙文',
        'ru': '俄文',
        'it': '意大利文',
        'pt': '葡萄牙文',
    }

    def __init__(self, model_name: str = "base"):
        """
        初始化Whisper实现
        
        Args:
            model_name: Whisper模型名称，默认为"base"
        """
        # 初始化日志记录器
        self.logger = get_logger(__name__)
        # 设置模型名称
        self.model_name = model_name
        # 初始化模型实例
        self.model = None
        # 加载Whisper模型
        self._load_model()

    def _load_model(self):
        """
        加载Whisper模型
        """
        try:
            # 使用whisper库加载指定模型
            self.model = whisper.load_model(self.model_name)
        except Exception as e:
            self.logger.error(f"加载Whisper模型失败: {e}")
            raise e

    def speech_to_text(self, audio_file: str, language: str = None, **kwargs) -> Tuple[BaseCodeEnum, Dict[str, Any]]:
        """
        使用Whisper将语音转换为文本
        
        Args:
            audio_file: 音频文件路径
            language: 语言代码，如果为None则自动检测
            **kwargs: 扩展参数
            
        Returns:
            Tuple[int, str, Dict[str, Any]]: (状态码, 识别文本, 结果信息)
        """
        self.logger.info(f"开始语音转文本: audio_file={audio_file}, language={language}")

        try:
            # 准备转录参数
            transcribe_kwargs = {}
            if language:
                transcribe_kwargs["language"] = language

            # 执行语音转文本
            result = self.model.transcribe(audio_file, **transcribe_kwargs)

            # 提取识别结果
            recognized_text = result.get("text", "").strip()
            detected_language = result.get("language", "unknown")

            self.logger.info(f"语音转文本完成: 识别文本={recognized_text}, 语言={detected_language}")
            # 构建结果信息
            result_info = {
                "text": recognized_text,
                "language": detected_language,
                "language_confidence": result.get("language_probs", {}).get(detected_language, 0.0),
                "segments": len(result.get("segments", [])),
                "duration": result.get("segments", [{}])[-1].get("end", 0.0) if result.get("segments") else 0.0
            }

            return ReturnCode.SUCCESS, result_info

        except Exception as e:
            self.logger.error(f"语音转文本失败: {e}")
            raise e

    def get_supported_languages(self) -> Dict[str, str]:
        """
        获取支持的语言列表
        
        Returns:
            Dict[str, str]: 语言代码到语言名称的映射
        """
        # 返回支持的语言映射副本
        return self.SUPPORTED_LANGUAGES.copy()

    def get_model_info(self) -> Dict[str, Any]:
        """
        获取模型信息
        
        Returns:
            Dict[str, Any]: 包含模型名称、版本、大小等信息的字典
        """
        # 检查模型是否已加载
        if self.model is None:
            return {"error": "模型未加载"}

        # 返回模型详细信息
        return {
            "model_name": self.model_name,
            "model_type": "whisper",
            "version": "1.0.0",
            "size_mb": self.model.dims.n_audio_state * 4 / 1024 / 1024 if hasattr(self.model, 'dims') else 0
        }
